# coding: utf-8
__author__ = 't_wolf'

import re
import paramiko
from os import linesep


class SSHError(BaseException):
    def __init__(self):
        super(SSHError, self).__init__(self)


class SSH:
    #
    header_buffer, command_sign = None, None
    std_in, std_out, std_err = None, None, None

    def __init__(self, dd=None, end_expect="#"):
        self.end_expect = end_expect
        self.__ssh = paramiko.SSHClient()
        self.__ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        if dd is not None:
            self.connect(dd)

    def connect(self, dd):
        try:
            self.__ssh.connect(dd['IP'], dd['PORT'], dd['UserName'], dd['Password'])
            self.__shell = self.__ssh.invoke_shell()
            self.header_buffer = self.get_receive().strip() + " "
            self.command_sign = self.header_buffer.split("\n")[-1]
            self.__transport = paramiko.Transport((dd['IP'], dd['PORT']))
            self.__transport.connect(username=dd['UserName'], password=dd['Password'])
            self.__sftp = paramiko.SFTPClient.from_transport(self.__transport)
        except Exception, e:
            print e

    def execute(self, command):
        self.std_in, self.std_out, self.std_err = \
            self.__ssh.exec_command(command)
        return self.std_in, self.std_out, self.std_err

    def get_receive(self, _end_expect=None):
        __end_expect = _end_expect if _end_expect != None else self.end_expect
        buff = ""
        while not buff.strip().endswith(__end_expect):
            buff += self.__shell.recv(1024)
        # _res = buff.split(linesep)
        return re.sub("\[.*?m", "", buff.replace(chr(13), "").replace("\x1b", ""))
        # return check_code(buff)
        # if len(_res) <= 2: return None
        # else:
        #     del _res[0], _res[-1]
        #     return (_res, buff)

    def communicate(self, _cmd, _end_expect=None):
        if self.header_buffer != None:
            buff = self.header_buffer
            self.header_buffer = None
        else:
            buff = self.command_sign

        self.__shell.send(_cmd + "\n")
        _temp_buff = self.get_receive()
        buff += _temp_buff[:_temp_buff.rindex("\n")]

        self.__shell.send("echo $?\n")
        code = int(self.get_receive().split("\n")[1])
        return code, buff

    def get(self, remoteName, localName, callback=""):
        try:
            self.__sftp.get(remoteName, localName, callback)
        except IOError:
            print "Must give absolutely Path + FileName."
        return 0

    def put(self, localName, remoteName, callback=""):
        try:
            self.__sftp.put(localName, remoteName, callback)
        except IOError:
            print "Must give absolutely Path + FileName."
        return 0

    def close(self):
        try:
            self.__ssh.close()
            self.__transport.close()
        except Exception:
            return 128
        else:
            return 0

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()


if __name__ == '__main__':
    # com_list = [
    #     "ps -ef | grep ls",
    #     "ps -ef | grep proxy-web | grep -v grep | awk '{print $2}'",
    #     "ls /root"
    # ]
    # with SSH({
    #     "IP": "10.21.16.32",
    #     "PORT": 22,
    #     "UserName": "bds-agent",
    #     "Password": "asiainfo"
    #     #"Password": "as!a1nFo#!eNc_"
    # }, '$') as ssh:
    #     for i in com_list:
    #         print ssh.communicate(i)[1]
    print SSH({
        "IP": "192.168.188.210",
        "PORT": 22,
        "UserName": "admin",
        "Password": "1qazXSW@3edc"
    }, '#').communicate("echo whatfuck | passwd --stdin root")
